package indi.madoka.weeb.module.base.service;

import com.theokanning.openai.completion.CompletionChoice;
import com.theokanning.openai.completion.CompletionRequest;
import com.theokanning.openai.image.CreateImageRequest;
import com.theokanning.openai.image.Image;
import com.theokanning.openai.model.Model;
import com.theokanning.openai.service.OpenAiService;
import indi.madoka.weeb.utils.RedisOperator;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.springframework.stereotype.Service;

import java.time.Duration;
import java.util.*;

@Service
@Slf4j
public class CustomOpenAiService {

    private static final String TOKEN = "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";

    private static final String NO_REPLY = "我不太理解，尝试换个话题或者清除语境";

    private static final String TOO_LENGTHY_CONTEXT = "当前ID：{%s}语境累计过长，已重置";

    private static final String SINGLE_KEY_PREFIX = "single_context_key:";

    private static final String GROUP_KEY_PREFIX = "group_context_key:";

    private static final String MODEL = "text-davinci-003";

    private static final Integer MAX_TOKEN = 2048;

    private static final List<String> STOP = Collections.singletonList("Human:");

    private final OpenAiService openAiService;

    private final RedisOperator redisOperator;

    public CustomOpenAiService(RedisOperator redisOperator) {
        this.openAiService = new OpenAiService(TOKEN, Duration.ofMinutes(1));
        this.redisOperator = redisOperator;
    }

    private String getChatReply(String prompt, Integer id, String prefix) {
        try {
            String context = this.getContext(prefix, id, prompt);
            if(StringUtils.isEmpty(context))
                return String.format(TOO_LENGTHY_CONTEXT, id);
            CompletionRequest completionRequest = CompletionRequest.builder()
                    .prompt(context)
                    .model(MODEL)
                    .maxTokens(MAX_TOKEN)
                    .temperature(0.5)
                    .stop(STOP)
                    .build();
            List<CompletionChoice> completionChoices = Optional.ofNullable(
                    openAiService.createCompletion(completionRequest).getChoices())
                    .orElseGet(ArrayList::new);
            String text;
            if(completionChoices.isEmpty() || StringUtils.isEmpty(text = completionChoices.get(0).getText())) {
                this.clearLastPrompt(prefix, id);
                return NO_REPLY;
            }
            String reply = this.getPrettyText(text);
            this.setContext(prefix, id, context, reply);
            return reply;
        }catch (Exception e) {
            this.clearLastPrompt(prefix, id);
            log.error("[OP-LOG]-[OpenService]-getChatReply({}):{}", prompt, e.getMessage());
            return "执行错误，已删除上一次对话，再次尝试或清除语境";
        }
    }

    public String getChatReply(String prompt, String model) {
        try {
            CompletionRequest completionRequest = CompletionRequest.builder()
                    .prompt(prompt)
                    .model(model)
                    .maxTokens(MAX_TOKEN)
                    .temperature(0.5)
                    .echo(true)
                    .build();
            List<CompletionChoice> completionChoices = Optional.ofNullable(
                    openAiService.createCompletion(completionRequest).getChoices()).orElseGet(ArrayList::new);
            String text;
            if(completionChoices.isEmpty() || Objects.isNull(completionChoices.get(0))
                    || StringUtils.isEmpty(text = completionChoices.get(0).getText()))
                return NO_REPLY;
            return text;
        } catch (Exception e) {
            log.error("[OP-LOG]-[OpenService]-getChatReply({}, {}):{}", prompt, model, e.getMessage());
            return "执行错误，稍后再试";
        }
    }

    public String getCreateImage(String prompt) {
        try {
            CreateImageRequest  imageRequest = CreateImageRequest.builder()
                    .prompt(prompt)
                    .n(1)
                    .size("512x512")
                    .build();
            List<Image> images = Optional.ofNullable(
                    openAiService.createImage(imageRequest).getData()).orElseGet(ArrayList::new);
            String url;
            if(images.isEmpty() || StringUtils.isEmpty(url = images.get(0).getUrl()))
                return "绘图出错";
            return url;
        } catch (Exception e) {
            log.error("[OP-LOG]-[OpenService]-getCreateImage({}):{}", prompt, e.getMessage());
            return "执行错误，稍后再试";
        }
    }

    public String getModels() {
        try {
            return openAiService.listModels().stream()
                    .limit(20)
                    .map(Model::getId)
                    .reduce((x, y) -> x + ";" + y)
                    .orElse("获取模型列表失败");
        } catch (Exception e) {
            log.error("[OP-LOG]-[OpenService]-getModels():{}", e.getMessage());
            return "执行错误，稍后再试";
        }
    }

    public String getGroupChatReply(String prompt, Integer gid) {
        return this.getChatReply(prompt, gid, GROUP_KEY_PREFIX);
    }

    public String getSingleChatReply(String prompt, Integer id) {
        return this.getChatReply(prompt, id, SINGLE_KEY_PREFIX);
    }

    public String shareContext(Integer id, Integer gid) {
        String holder = redisOperator.get(SINGLE_KEY_PREFIX + id);
        if(StringUtils.isEmpty(holder))
            return "当前ID：{ " + id +  "}尚没有建立语境";
        redisOperator.set(GROUP_KEY_PREFIX + gid, redisOperator.get(GROUP_KEY_PREFIX + gid) + "\n" + holder);
        return "已将当前语境分享到群聊";
    }

    public void clear(String prefix, Integer id) {
        redisOperator.del(prefix + id);
    }

    private String getContext(String prefix, Integer id, String prompt) {
        String holder = redisOperator.get(prefix + id);
        String context = (StringUtils.isEmpty(holder) ? "" : holder + "\n") + "Human:" + prompt;
        if(context.length() > (MAX_TOKEN << 3)) {
           this.clear(prefix, id);
            return null;
        }
        redisOperator.set(prefix + id, context);
        return context;
    }

    private void setContext(String prefix, Integer id, String context, String reply) {
        if(NO_REPLY.equals(reply)) {
         this.clearLastPrompt(prefix, id);
         return;
        }
        redisOperator.set(prefix + id, context + "\n" + reply);
    }

    private String getPrettyText(String text) {
        text = text.replaceFirst(".*?\\n{1,2}(.*?[:|：])?", "");
        return StringUtils.isBlank(text) ? NO_REPLY : text;
    }

    private void clearLastPrompt(String prefix, Integer id) {
        String context = redisOperator.get(prefix + id);
        redisOperator.set(prefix + id, context.substring(0, context.lastIndexOf(STOP.get(0))));
    }

    public static String getGroupKeyPrefix() {
        return GROUP_KEY_PREFIX;
    }

    public static String getSingleKeyPrefix() {
        return SINGLE_KEY_PREFIX;
    }
}
